import os
import pandas as pd
from collections import defaultdict
from openpyxl import Workbook
import tkinter as tk
from tkinter import filedialog, messagebox

def process_xlsx_files(folder_path, output_dir):
    """
    读取指定文件夹下所有 xlsx 文件，统计每年进出交易金额总数
    """
    # 初始化统计字典：{年份: {'进': 总金额, '出': 总金额}}
    yearly_stats = defaultdict(lambda: {'进': 0.0, '出': 0.0})
    
    # 获取文件夹下所有 .xlsx 文件
    xlsx_files = [f for f in os.listdir(folder_path) if f.endswith('.xlsx') and not f.startswith('.~') and not f.startswith('~$')]
    
    if not xlsx_files:
        print("文件夹中没有找到 .xlsx 文件")
        return yearly_stats
    
    print(f"找到 {len(xlsx_files)} 个 Excel 文件:")
    
    for file_name in xlsx_files:
        file_path = os.path.join(folder_path, file_name)
        print(f"  处理文件: {file_name}")
        
        try:
            # 读取 Excel 文件的所有 sheet
            xls = pd.ExcelFile(file_path)
            
            for sheet_name in xls.sheet_names:
                try:
                    df = pd.read_excel(file_path, sheet_name=sheet_name)
                    
                    # 检查必需的列是否存在
                    required_columns = ['借贷标识', '交易金额', '交易时间']
                    if not all(col in df.columns for col in required_columns):
                            print(f"    跳过 sheet '{sheet_name}'：缺少必需列")
                            continue
                                        
                    # 转换交易时间为 datetime 类型
                    df['交易时间'] = pd.to_datetime(df['交易时间'], errors='coerce')
                    
                    # 删除无效的时间数据
                    df = df.dropna(subset=['交易时间'])
                    # 展示df所有数据
                    # print(df)
                    # 遍历每一行数据
                    for _, row in df.iterrows():
                        flag = row['借贷标识']
                        money = row['交易金额']
                        trade_time = row['交易时间']
                        
                        # 检查数据有效性
                        if pd.isna(flag) or pd.isna(money) or not isinstance(money, (int, float)):
                            print(f"    跳过无效数据行: {row.to_dict()}")
                            continue
                        
                        year = trade_time.year
                        
                        if flag == '进':
                            yearly_stats[year]['进'] += float(money)
                        elif flag == '出':
                            yearly_stats[year]['出'] += float(money)

                    # 保存结果到Excel文件
                    save_results(yearly_stats, file_name, sheet_name, output_dir)
                    # display_results(yearly_stats, file_name, sheet_name)  # debug使用

                    # 清空统计数据以便处理下一个sheet
                    yearly_stats.clear()     
                except Exception as e:
                    print(f"    处理 sheet '{sheet_name}' 时出错: {str(e)}")
                    continue
                    
        except Exception as e:
            print(f"  处理文件 '{file_name}' 时出错: {str(e)}")
            continue
    
    return yearly_stats

def save_results(yearly_stats, file_name, sheet_name, output_dir='数据处理结果'):
    """
    生成并保存统计结果到xlsx文件
    """
    if not yearly_stats:
        print("没有统计数据")
        return
    
    # 准备要写入Excel的数据
    data = []
    data.append(["年份", "进账总额(万元)", "出账总额(万元)"])
    
    # 按年份排序
    for year in sorted(yearly_stats.keys()):
        income = yearly_stats[year]['进']
        expense = yearly_stats[year]['出']
        data.append([year, round(income/10000.0, 2), round(expense/10000.0, 2)])
    
    # 计算总计
    total_income = sum(stats['进'] for stats in yearly_stats.values())
    total_expense = sum(stats['出'] for stats in yearly_stats.values())
    data.append(["总计", round(total_income/10000.0, 2), round(total_expense/10000.0, 2)])
   
    # 创建工作簿和工作表
    wb = Workbook()
    ws = wb.active
    ws.title = f"{sheet_name[:31]}"  # Excel工作表名称限制为31个字符
    
    # 写入数据
    for row in data:
        ws.append(row)
    
    # 保存文件
    output_filename = os.path.join(output_dir, file_name)
    wb.save(output_filename)
    print(f"文件已保存到 {output_filename}")

def display_results(yearly_stats, file_name, sheet_name):
    """
    显示统计结果
    """
    if not yearly_stats:
        print("没有统计数据")
        return
    
    print("\n" + "="*50)
    print(f"{file_name} {sheet_name} 每年进出交易金额统计结果")
    print("="*50)
    print(f"{'年份':<8} {'进账总额':<15} {'出账总额':<15}")
    print("-"*50)
    
    # 按年份排序
    for year in sorted(yearly_stats.keys()):
        income = yearly_stats[year]['进']
        expense = yearly_stats[year]['出']
        print(f"{year:<8} {income/10000.0:<15.2f} {expense/10000.0:<15.2f}")
    
    print("-"*50)
    # 计算总计
    total_income = sum(stats['进'] for stats in yearly_stats.values())
    total_expense = sum(stats['出'] for stats in yearly_stats.values())
    print(f"{'总计':<8} {total_income/10000.0:<15.2f} {total_expense/10000.0:<15.2f}")

def select_folder():
    """
    创建一个简单的UI界面用于选择文件夹路径
    """
    root = tk.Tk()
    root.withdraw()  # 隐藏主窗口
    
    # 打开文件夹选择对话框
    folder_path = filedialog.askdirectory(title="请选择需要处理的文件夹")
    
    # 销毁根窗口
    root.destroy()
    
    return folder_path

def main():
    """
    主函数：通过UI选择文件夹并处理Excel文件
    """
    print("请选择需要处理的文件夹...")
    
    # 通过UI选择文件夹
    folder_path = select_folder()
    
    if not folder_path:
        print("未选择文件夹")
        return
    
    if not os.path.exists(folder_path):
        print(f"选择的文件夹不存在: {folder_path}")
        return
    
    print(f"选择的文件夹: {folder_path}")
    
    output_dir = '数据处理结果'
    # 保存结果的文件夹，如果没有就新建
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    # 处理文件
    stats = process_xlsx_files(folder_path, output_dir)

if __name__ == "__main__":
    main()
    #按任意键退出程序
    input("按任意键退出程序...")